#!/usr/bin/env bash
# BERT 分类器训练脚本
# 
# 使用示例:
#   ./run_train_eval.sh                                    # 使用默认配置
#   DATA_DIR=/path/to/data ./run_train_eval.sh             # 自定义数据目录
#   MODEL_NAME=bert-base-chinese EPOCHS=5 ./run_train_eval.sh  # 自定义模型和训练轮数
#
# 环境变量配置:
#   DATA_DIR      - 训练数据目录 (默认: bert_train/data)
#   MODEL_DIR     - 模型输出目录 (默认: bert_train/models) 
#   RESULT_DIR    - 评测结果目录 (默认: ${DATA_DIR}/result)
#   INPUT_PATH    - 输入数据文件路径 (默认: 自动查找最新的auto_pass_*.jsonl)
#   MODEL_NAME    - 预训练模型名称 (默认: ckiplab/bert-tiny-chinese)
#   EPOCHS        - 训练轮数 (默认: 3)
#   BATCH_SIZE    - 批次大小 (默认: 16)
#   MAX_LENGTH    - 最大序列长度 (默认: 512)
#   TEXT_FORMAT   - 文本格式 (默认: bert_pair, 可选: with_context)
#   FILTER_BY_TOKENS - 是否按token长度过滤 (默认: 0, 设为1启用)

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"

# 可配置参数
DATA_DIR="${DATA_DIR:-${ROOT_DIR}/bert_train/data}"
MODEL_DIR="${MODEL_DIR:-${ROOT_DIR}/bert_train/models}"
RESULT_DIR="${RESULT_DIR:-${DATA_DIR}/result}"
MODEL_NAME="${MODEL_NAME:-ckiplab/bert-tiny-chinese}"
EPOCHS=${EPOCHS:-3}
BATCH_SIZE=${BATCH_SIZE:-16}
MAX_LENGTH=${MAX_LENGTH:-512}
TEXT_FORMAT="${TEXT_FORMAT:-bert_pair}"
INPUT_PATH="${INPUT_PATH:-${ROOT_DIR}/data/label_data/auto_pass.jsonl}"

echo "[1/3] 检查数据... (数据目录: ${DATA_DIR}, 文本格式: ${TEXT_FORMAT})"
NEED_PREP=0
for i in 1 2 3 4 5; do
  if [ ! -s "${DATA_DIR}/fold_${i}.jsonl" ]; then
    NEED_PREP=1
    break
  fi
done
if [ ! -s "${DATA_DIR}/split_meta.json" ]; then
  NEED_PREP=1
fi

if [ ${NEED_PREP} -eq 1 ]; then
  echo "未检测到完整折分数据，开始生成..."
  mkdir -p "${DATA_DIR}"
  PREP_ARGS=("--output_dir" "${DATA_DIR}" "--text_format" "${TEXT_FORMAT}")
  if [ "${FILTER_BY_TOKENS:-0}" = "1" ]; then
    PREP_ARGS+=("--filter_by_tokens")
  fi
  if [ -n "${INPUT_PATH:-}" ]; then
    PREP_ARGS+=("--input" "${INPUT_PATH}")
  fi
  python3 "${ROOT_DIR}/bert_train/prepare_core_dataset.py" "${PREP_ARGS[@]}"
else
  echo "检测到现有折分数据，跳过生成步骤。"
fi

echo "[2/3] 训练模型（fold_1-4为训练集，fold_5不参与）..."
python3 "${ROOT_DIR}/bert_train/train_classifier.py" \
  --data_dir "${DATA_DIR}" \
  --output_dir "${MODEL_DIR}" \
  --model_name "${MODEL_NAME}" \
  --epochs "${EPOCHS}" \
  --batch_size "${BATCH_SIZE}" \
  --max_length "${MAX_LENGTH}"

# 获取最新的checkpoint路径
LATEST_MODEL_DIR=$(python3 -c "
import glob
checkpoints = glob.glob('${MODEL_DIR}/checkpoint-*')
if checkpoints:
    latest = max(checkpoints, key=lambda x: int(x.split('-')[-1]) if x.split('-')[-1].isdigit() else 0)
    print(latest)
else:
    print('${MODEL_DIR}')
")

echo "[3/3] 评测模型（使用fold_5作为测试集）..."
python3 "${ROOT_DIR}/bert_train/eval_classifier.py" \
  --data_dir "${DATA_DIR}" \
  --model_dir "${LATEST_MODEL_DIR}" \
  --result_dir "${RESULT_DIR}" \
  --max_length "${MAX_LENGTH}"

echo "完成。使用模型: ${LATEST_MODEL_DIR}"
echo "训练checkpoints: ${MODEL_DIR}"
echo "评测结果: ${RESULT_DIR}"


